{
    "nbformat": 4,
    "nbformat_minor": 4,
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "version": 3,
                "name": "ipython"
            },
            "mimetype": "text/x-python",
            "file_extension": ".py",
            "version": "3.6.3",
            "pygments_lexer": "ipython3",
            "nbconvert_exporter": "python",
            "name": "python"
        }
    },
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# [HW5] Problem: Total Least Squares\n",
                "\n",
                "Import necessary Python packages.\n"
            ]
        },
        {
            "execution_count": null,
            "cell_type": "code",
            "metadata": {},
            "source": [
                "import numpy as np\n",
                "import matplotlib.pyplot as plt\n",
                "from imageio import imread\n",
                "from imageio import imwrite as imsave\n",
                "import os\n",
                "\n",
                "data_dir = \"imgs\"\n"
            ],
            "outputs": []
        },
        {
            "execution_count": null,
            "cell_type": "code",
            "metadata": {},
            "source": [
                "# This loads and returns all of the images needed for the problem\n",
                "# data - the image of the spherical mirror\n",
                "# tennis - the image of the tennis ball that we will relight\n",
                "# target - the image that we will paste the tennis ball onto\n",
                "def loadImages():\n",
                "    imFile = 'stpeters_probe_small.png'\n",
                "    compositeFile = 'tennis.png'\n",
                "    targetFile = 'interior.jpg'\n",
                "\n",
                "    data = imread( os.path.join(data_dir, imFile) ).astype('float')*1.5\n",
                "    tennis = imread( os.path.join(data_dir, compositeFile)).astype('float')\n",
                "    target = imread( os.path.join(data_dir, targetFile)).astype('float')/255\n",
                "\n",
                "    return data, tennis, target\n",
                "\n",
                "\n",
                "# This function takes as input a square image of size m x m x c\n",
                "# where c is the number of color channels in the image.  We\n",
                "# assume that the image contains a scphere and that the edges\n",
                "# of the sphere touch the edge of the image.\n",
                "# The output is a tuple (ns, vs) where ns is an n x 3 matrix\n",
                "# where each row is a unit vector of the direction of incoming light\n",
                "# vs is an n x c vector where the ith row corresponds with the\n",
                "# image intensity of incoming light from the corresponding row in ns\n",
                "def extractNormals(img):\n",
                "\n",
                "    # Assumes the image is square\n",
                "    d = img.shape[0]\n",
                "    r = d / 2\n",
                "    ns = []\n",
                "    vs = []\n",
                "    for i in range(d):\n",
                "        for j in range(d):\n",
                "\n",
                "            # Determine if the pixel is on the sphere\n",
                "            x = j - r\n",
                "            y = i - r\n",
                "            if x*x + y*y > r*r-100:\n",
                "                continue\n",
                "\n",
                "            # Figure out the normal vector at the point\n",
                "            # We assume that the image is an orthographic projection\n",
                "            z = np.sqrt(r*r-x*x-y*y)\n",
                "            n = np.asarray([x,y,z])\n",
                "            n = n / np.sqrt(np.sum(np.square(n)))\n",
                "            view = np.asarray([0,0,-1])\n",
                "            n = 2*n*(np.sum(n*view))-view\n",
                "            ns.append(n)\n",
                "            vs.append(img[i,j])\n",
                "\n",
                "    return np.asarray(ns), np.asarray(vs)\n",
                "\n",
                "# This function renders a diffuse sphere of radius r\n",
                "# using the spherical harmonic coefficients given in\n",
                "# the input coeff where coeff is a 9 x c matrix\n",
                "# with c being the number of color channels\n",
                "# The output is an 2r x 2r x c image of a diffuse sphere\n",
                "# and the value of -1 on the image where there is no sphere\n",
                "def renderSphere(r,coeff):\n",
                "\n",
                "    d = 2*r\n",
                "    img = -np.ones((d,d,3))\n",
                "    ns = []\n",
                "    ps = []\n",
                "\n",
                "    for i in range(d):\n",
                "        for j in range(d):\n",
                "\n",
                "            # Determine if the pixel is on the sphere\n",
                "            x = j - r\n",
                "            y = i - r\n",
                "            if x*x + y*y > r*r:\n",
                "                continue\n",
                "\n",
                "            # Figure out the normal vector at the point\n",
                "            # We assume that the image is an orthographic projection\n",
                "            z = np.sqrt(r*r-x*x-y*y)\n",
                "            n = np.asarray([x,y,z])\n",
                "            n = n / np.sqrt(np.sum(np.square(n)))\n",
                "            ns.append(n)\n",
                "            ps.append((i,j))\n",
                "\n",
                "    ns = np.asarray(ns)\n",
                "    B = computeBasis(ns)\n",
                "    vs = B.dot(coeff)\n",
                "\n",
                "    for p,v in zip(ps,vs):\n",
                "        img[p[0],p[1]] = np.clip(v,0,255)\n",
                "\n",
                "    return img\n",
                "\n",
                "# relights the sphere in img, which is assumed to be a square image\n",
                "# coeff is the matrix of spherical harmonic coefficients\n",
                "def relightSphere(img, coeff):\n",
                "    img = renderSphere(int(img.shape[0]/2),coeff)/255*img/255\n",
                "    return img\n",
                "\n",
                "# Copies the image of source onto target\n",
                "# pixels with values of -1 in source will not be copied\n",
                "def compositeImages(source, target):\n",
                "\n",
                "    # Assumes that all pixels not equal to 0 should be copied\n",
                "    out = target.copy()\n",
                "    cx = int(target.shape[1]/2)\n",
                "    cy = int(target.shape[0]/2)\n",
                "    sx = cx - int(source.shape[1]/2)\n",
                "    sy = cy - int(source.shape[0]/2)\n",
                "\n",
                "    for i in range(source.shape[0]):\n",
                "        for j in range(source.shape[1]):\n",
                "            if np.sum(source[i,j]) >= 0:\n",
                "                out[sy+i,sx+j] = source[i,j]\n",
                "\n",
                "    return out\n",
                "\n",
                "# Fill in this function to compute the basis functions\n",
                "# This function is used in renderSphere()\n",
                "def computeBasis(ns):\n",
                "    # Returns the first 9 spherical harmonic basis functions\n",
                "\n",
                "    #################################################\n",
                "    # TODO: Compute the first 9 basis functions [Part (e)]\n",
                "    #################################################\n",
                "    B = np.ones((len(ns),9)) # This line is here just to fill space\n",
                "\n",
                "    ### start 1 ###\n",
                "\n",
                "    ### end 1 ###\n",
                "    return B\n",
                "\n",
                "if __name__ == '__main__':\n",
                "\n",
                "    data,tennis,target = loadImages()\n",
                "    ns, vs = extractNormals(data)\n",
                "    B = computeBasis(ns)\n",
                "\n",
                "    # reduce the number of samples because computing the SVD on\n",
                "    # the entire data set takes too long\n",
                "    Bp = B[::50]\n",
                "    vsp = vs[::50]\n",
                "\n",
                "    #################################################\n",
                "    # TODO: Solve for the coefficients using least squares  [Part (e)],\n",
                "    # total least squares  [Part (f)] or total least squares [Part (g)] here\n",
                "    # Below is a starter code\n",
                "    ##################################################\n",
                "    coeff = np.zeros((9,3))\n",
                "    coeff[0,:] = 255\n",
                "\n",
                "    img = relightSphere(tennis,coeff)\n",
                "\n",
                "    output = compositeImages(img,target)\n",
                "\n",
                "    print('Coefficients:\\n'+str(coeff))\n",
                "\n",
                "    plt.figure(1)\n",
                "    plt.imshow(output)\n",
                "    plt.show()\n",
                "\n",
                "    imsave( 'output.png',output)\n",
                "\n",
                "    ### start 2 ###\n",
                "\n",
                "    ### end 2 ###\n"
            ],
            "outputs": []
        }
    ]
}